# -*- coding: utf-8 -*-
import numpy as np
import torch
from torch import nn

class lstm_net(nn.Module):
    def __init__(self, 
        input_size:int, hidden_size:int, num_layers:int, out_dim:int):
        super(lstm_net, self).__init__()
        self.lstm = nn.LSTM(input_size=input_size, 
                    hidden_size=hidden_size,
                    num_layers=num_layers, 
                    batch_first=True,)
        self.dense = nn.Sequential(nn.Linear(self.lstm.hidden_size, out_dim), 
                                    nn.Sigmoid())
    def forward(self, X, state=None):
        X, self.state = self.lstm(X, state)
        Y = self.dense(X)
        return(Y, self.state)

if __name__ == "__main__":
    model = lstm_net(input_size=200,
                    hidden_size=256, 
                    num_layers=2,
                    out_dim=6)
    # input size (batch_size, length, features)
    a = torch.zeros(10, 3, 200)
    state = None
    b, state = model(a, state)
    print(a.size())
    # return sequence
    print(b.size())
    
